#!/usr/bin/env python3

from sys import exit
import argparse

version = '1.0 (2016-02-24)'

# Set up & parse the command line arguments:
parser = argparse.ArgumentParser(description = 'Basic FASTQ manipulation in Python', fromfile_prefix_chars='@')
parser.add_argument('-v', '--version', action='version', version='%(prog)s {0}'.format(version))
parser.add_argument(dest='command', choices=['view', 'fasta', 'header', 'sequence', 'header2', 'quality', 'basetab', 'qualtab', 'lengthtab'], default='view', help='The command to run')
parser.add_argument(dest='filename', help='The FASTQ file to process')
args = parser.parse_args()

# Function to return an error:
def error(message): exit('ERROR: {}'.format(message))

class FQRead(object):
    total_basechars = '''ABCDGHKMNRSTUVWXYabcdghkmnrstuvwxy'''
    total_qualchars = '''!"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`abcdefghijklmnopqrstuvwxyz{|}~'''
    def __init__(self, f):
        self.header = f.readline().strip().lstrip('@')
        if self.header == '': raise Exception()
        self.sequence = f.readline().strip()
        self.header2 = f.readline().strip().lstrip('+')
        self.quality = f.readline().strip()
    def __len__(self): return len(self.sequence)
    def printFASTQ(self):
        print('@{}'.format(self.header))
        print(self.sequence)
        print('+{}'.format(self.header2))
        print(self.quality)
    def printFASTA(self):
        print('>{}'.format(self.header))
        print(self.sequence)
    def baseTab(self):
        output = {}
        for base in FQRead.total_basechars:
            output[base] = self.sequence.count(base)
        return output
    def qualTab(self):
        output = {}
        for base in FQRead.total_qualchars:
            output[base] = self.quality.count(base)
        return output

# Open the input file:
try: input_file = open(args.filename, 'rt')
except: error('failed to open file "{}"'.format(args.filename))

if args.command == 'fasta':
    try:
        while True: FQRead(input_file).printFASTA()
    except: pass

elif args.command == 'header':
    try:
        while True:
            print(FQRead(input_file).header)
    except: pass
    
elif args.command == 'sequence':
    try:
        while True:
            print(FQRead(input_file).sequence)
    except: pass
    
elif args.command == 'header2':
    try:
        while True:
            print(FQRead(input_file).header2)
    except: pass
    
elif args.command == 'quality':
    try:
        while True:
            print(FQRead(input_file).quality)
    except: pass

elif args.command == 'basetab':
    totals = {}
    for base in FQRead.total_basechars: totals[base] = 0
    try:
        while True:
            basetab = FQRead(input_file).baseTab()
            for base in FQRead.total_basechars: totals[base] += basetab[base]            
    except: pass
    for base in FQRead.total_basechars: print('{}\t{}'.format(base, totals[base]))

elif args.command == 'qualtab':
    totals = {}
    for base in FQRead.total_qualchars: totals[base] = 0
    try:
        while True:
            basetab = FQRead(input_file).qualTab()
            for base in FQRead.total_qualchars: totals[base] += basetab[base]            
    except: pass
    for base in FQRead.total_qualchars: print('{}\t{}'.format(base, totals[base]))

elif args.command == 'lengthtab':
    totals = {}
    try:
        while True:
            length = len(FQRead(input_file))
            if length not in totals.keys(): totals[length] = 1
            else: totals[length] += 1
    except: pass
    for length in sorted(totals.keys()):
        print('{}\t{}'.format(length, totals[length]))

else:
    try:
        while True: FQRead(input_file).printFASTQ()
    except: pass

